# uncompyle6 version 3.2.3
# Python bytecode 3.6 (3379)
# Decompiled from: Python 3.6.8 |Anaconda custom (64-bit)| (default, Feb 21 2019, 18:30:04) [MSC v.1916 64 bit (AMD64)]
# Embedded file name: api\handlers\frontend\workers.py
__author__ = "Attila Gerendi (Sanyi)"
import logging, base64, os.path
from helpers.licensing import get_local_license
from helpers.hashing import md5
from sqlalchemy import select as sql_select
from api.handlers.frontend import FrontendResourceHandler
from api.classes.frontend.auto import *
from server.http_responses.http_status_codes import Http204
from server.http_responses.http_errors import Http409, Http404, Http403
from db.data.users import (
    require_system_user,
    require_at_least_tech_admin,
    is_at_least_tech_admin,
)
from helpers.setup_types import is_setup_type_on_premise_master
from db import Connection
from db.tables.workers import (
    WorkersTable,
    WorkerRow,
    WorkerAuthorization,
    WorkerStatus,
    LocalWorker,
)
from db.tables.targets_configurations import (
    TargetConfigurationRow,
    TargetConfigurationTable,
)
from db.data.access import generate_target_query, apply_filters
from db.tables.targets import TargetRow
from helpers.http_clients.scan_api import ScanApiClient
from helpers.http_clients.exceptions import *
from settings_provider import settings
from helpers.licensing.features import Features
from sqlalchemy.sql.functions import count as sql_count

logger = logging.getLogger("frontend_api_resources.Status")


class WorkerManager(FrontendResourceHandler):
    def get_workers_assigned_to_target(self, target_id, **_):
        if not is_setup_type_on_premise_master():
            raise Http404()
        self.license.require(Features.MULTI_ENGINE)
        if not is_at_least_tech_admin(self.request_user):
            return ClsWorkerList(workers=[])
        with Connection(self.current_shard_name) as (db):
            select_from, filters = generate_target_query(self.request_user)
            q = (
                sql_select((TargetRow.target_id,))
                .select_from(select_from)
                .where(TargetRow.target_id == target_id)
            )
            q = apply_filters(q, filters)
            if not db.execute(q).scalar():
                raise Http404()
            q = (
                sql_select((TargetConfigurationRow.value,))
                .where(TargetConfigurationRow.target_id == target_id)
                .where(TargetConfigurationRow.name == "workers")
            )
            target_workers = db.execute(q).scalar()
            if not target_workers:
                return ClsWorkerList(workers=[])
            q = (
                sql_select((WorkerRow,))
                .where(WorkerRow.worker_id.in_(target_workers.values()))
                .order_by(
                    WorkerRow.description, WorkerRow.endpoint, WorkerRow.worker_id
                )
            )
            workers = list()
            if LocalWorker.WORKER_ID in target_workers.values():
                workers.append(
                    ClsWorker(
                        scanning_app=LocalWorker.SCANNING_APP,
                        endpoint=LocalWorker.DESCRIPTION,
                        description=LocalWorker.DESCRIPTION,
                        worker_id=LocalWorker.WORKER_ID,
                        status=WorkerStatus.ONLINE,
                        authorization=WorkerAuthorization.AUTHORIZED,
                        app_version=settings.get("version_info", {}).get("str", ""),
                    )
                )
            for worker in db.execute(q).fetchall():
                try:
                    app_version = worker.info.get("app_version", "")
                except:
                    app_version = ""

                workers.append(
                    ClsWorker(
                        scanning_app=worker.scanning_app,
                        endpoint=worker.endpoint,
                        description=worker.description,
                        worker_id=worker.worker_id,
                        status=worker.status,
                        authorization=worker.authorization,
                        app_version=app_version,
                    )
                )

            return ClsWorkerList(workers=workers)

    def assign_workers_to_target(self, target_id, body, **_):
        if not is_setup_type_on_premise_master():
            raise Http404()
        self.license.require(Features.MULTI_ENGINE)
        require_at_least_tech_admin(self.request_user)
        worker_id_list = (ClsWorkerIdList(**body)).get("worker_id_list")
        with Connection(self.current_shard_name) as (db):
            select_from, filters = generate_target_query(self.request_user)
            q = (
                sql_select((TargetRow.target_id,))
                .select_from(select_from)
                .where(TargetRow.target_id == target_id)
            )
            q = apply_filters(q, filters)
            if not db.execute(q).scalar():
                raise Http404()
            if not worker_id_list:
                q = (
                    TargetConfigurationTable.delete()
                    .where(TargetConfigurationRow.target_id == target_id)
                    .where(TargetConfigurationRow.name == "workers")
                )
                db.execute(q)
                return Http204()
            q = sql_select((WorkerRow,)).where(WorkerRow.worker_id.in_(worker_id_list))
            available_workers = dict()
            if LocalWorker.WORKER_ID in worker_id_list:
                available_workers[LocalWorker.SCANNING_APP] = LocalWorker.WORKER_ID
            for worker in db.execute(q).fetchall():
                if worker.authorization != WorkerAuthorization.AUTHORIZED:
                    raise Http409("invalid worker status")
                if worker.scanning_app in available_workers:
                    raise Http409("duplicate scanning app")
                available_workers[worker.scanning_app] = worker.worker_id

            for worker_id in worker_id_list:
                if worker_id == LocalWorker.WORKER_ID:
                    pass
                else:
                    if worker_id not in available_workers.values():
                        raise Http409("worker not found")

            q = (
                TargetConfigurationTable.delete()
                .where(TargetConfigurationRow.target_id == target_id)
                .where(TargetConfigurationRow.name == "workers")
            )
            db.execute(q)
            q = TargetConfigurationTable.insert().values(
                target_id=target_id, name="workers", value=available_workers
            )
            db.execute(q)
            return Http204()

    def get_workers(self, **_):
        if not is_setup_type_on_premise_master():
            raise Http404()
        self.license.require(Features.MULTI_ENGINE)
        require_system_user(self.request_user)
        workers = list()
        workers.append(
            ClsWorker(
                scanning_app=LocalWorker.SCANNING_APP,
                endpoint=LocalWorker.DESCRIPTION,
                description=LocalWorker.DESCRIPTION,
                worker_id=LocalWorker.WORKER_ID,
                status=WorkerStatus.ONLINE,
                authorization=WorkerAuthorization.AUTHORIZED,
                app_version=settings.get("version_info", {}).get("str", ""),
            )
        )
        with Connection(self.current_shard_name) as (db):
            q = sql_select((WorkerRow,)).order_by(
                WorkerRow.description, WorkerRow.endpoint, WorkerRow.worker_id
            )
            for worker in db.execute(q).fetchall():
                try:
                    app_version = worker.info.get("app_version", "")
                except:
                    app_version = ""

                workers.append(
                    ClsWorker(
                        scanning_app=worker.scanning_app,
                        endpoint=worker.endpoint,
                        description=worker.description,
                        worker_id=worker.worker_id,
                        status=worker.status,
                        authorization=worker.authorization,
                        app_version=app_version,
                    )
                )

        return ClsWorkerList(workers=workers)

    def get_worker(self, worker_id, **_):
        if not is_setup_type_on_premise_master():
            raise Http404()
        self.license.require(Features.MULTI_ENGINE)
        require_system_user(self.request_user)
        if worker_id == LocalWorker.WORKER_ID:
            return ClsWorker(
                scanning_app=LocalWorker.SCANNING_APP,
                endpoint="",
                description=LocalWorker.DESCRIPTION,
                worker_id=LocalWorker.WORKER_ID,
                status=WorkerStatus.ONLINE,
                authorization=WorkerAuthorization.AUTHORIZED,
            )
        with Connection(self.current_shard_name) as (db):
            q = sql_select((WorkerRow,)).where(WorkerRow.worker_id == worker_id)
            worker = db.execute(q).fetchone()
            if not worker:
                raise Http404()
            return ClsWorker(
                scanning_app=worker.scanning_app,
                endpoint=worker.endpoint,
                description=worker.description,
                worker_id=worker.worker_id,
                status=worker.status,
                authorization=worker.authorization,
            )

    def rename_worker(self, worker_id, body, **_):
        if not is_setup_type_on_premise_master():
            raise Http404()
        self.license.require(Features.MULTI_ENGINE)
        require_system_user(self.request_user)
        if worker_id == LocalWorker.WORKER_ID:
            raise Http409("local worker operations not allowed")
        data = ClsWorkerDescription(**body)
        with Connection(self.current_shard_name) as (db):
            q = (WorkersTable.update(values=dict(description=data.description))).where(
                WorkerRow.worker_id == worker_id
            )
            db.execute(q)
        return Http204()

    def delete_worker(self, worker_id, **_):
        if not is_setup_type_on_premise_master():
            raise Http404()
        self.license.require(Features.MULTI_ENGINE)
        require_system_user(self.request_user)
        if worker_id == LocalWorker.WORKER_ID:
            raise Http409("local worker operations not allowed")
        with Connection(self.current_shard_name) as (db):
            q = sql_select((WorkerRow,)).where(WorkerRow.worker_id == worker_id)
            worker = db.execute(q).fetchone()
            if not worker:
                return Http204()
            q = (
                sql_select((TargetConfigurationRow.target_id,))
                .where(TargetConfigurationRow.name == "workers")
                .where(
                    TargetConfigurationRow.value[worker.scanning_app].astext
                    == worker_id
                )
            )
            target_id_list = []
            for r in db.execute(q).fetchall():
                target_id_list.append(r.target_id)

            if target_id_list:
                raise Http409("Worker is used", details=target_id_list)
            q = WorkersTable.delete().where(WorkerRow.worker_id == worker_id)
            db.execute(q)
            worker = dict(worker)
        try:
            client = ScanApiClient(
                worker["endpoint"],
                worker["scanning_app"],
                api_key=worker["api_key"],
                request_timeout=5,
                ssl_verify=worker["ssl_verify"],
            )
            client.delete_request("/worker/" + worker["worker_id"])
        except HTTPException as e:
            logging.exception("detaching worker failed with %s", e.result)
        except RequestError:
            logging.exception("detaching worker failed with request error")
        except Exception as e:
            logging.exception("detaching worker failed with %s", e)

        return Http204()

    def delete_worker_ignore_errors(self, worker_id, **_):
        if not is_setup_type_on_premise_master():
            raise Http404()
        self.license.require(Features.MULTI_ENGINE)
        require_system_user(self.request_user)
        if worker_id == LocalWorker.WORKER_ID:
            raise Http409("local worker operations not allowed")
        with Connection(self.current_shard_name) as (db):
            q = sql_select((WorkerRow,)).where(WorkerRow.worker_id == worker_id)
            worker = db.execute(q).fetchone()
            if not worker:
                return Http204()
            q = (
                TargetConfigurationTable.delete()
                .where(TargetConfigurationRow.name == "workers")
                .where(
                    TargetConfigurationRow.value[worker.scanning_app].astext
                    == worker_id
                )
            )
            db.execute(q)
            q = WorkersTable.delete().where(WorkerRow.worker_id == worker_id)
            db.execute(q)
            worker = dict(worker)
        try:
            client = ScanApiClient(
                worker["endpoint"],
                worker["scanning_app"],
                api_key=worker["api_key"],
                request_timeout=5,
                ssl_verify=worker["ssl_verify"],
            )
            client.delete_request("/worker/" + worker["worker_id"])
        except HTTPException as e:
            logging.exception("detaching worker failed with %s", e.result)
        except RequestError:
            logging.exception("detaching worker failed with request error")
        except Exception as e:
            logging.exception("detaching worker failed with %s", e)

        return Http204()

    def authorize_worker(self, worker_id, **_):
        if not is_setup_type_on_premise_master():
            raise Http404()
        self.license.require(Features.MULTI_ENGINE)
        license_key = get_local_license()["license_key"]
        if license_key:
            license_key_hash = md5(license_key)
        else:
            raise Http409("product is not activated")
        require_system_user(self.request_user)
        if worker_id == LocalWorker.WORKER_ID:
            raise Http409("local Scanning Engine operations are not allowed")
        if self.license.max_engines is None:
            max_engines = 100000
        else:
            max_engines = self.license.max_engines - 1
        if max_engines <= 0:
            raise Http403("too many scanning engines")
        with Connection(self.current_shard_name) as (db):
            q = sql_select((sql_count(WorkerRow.worker_id),)).where(
                WorkerRow.authorization == WorkerAuthorization.AUTHORIZED
            )
            c = db.execute(q).scalar()
            if c >= max_engines:
                raise Http403("too many scanning engines")
            q = sql_select((WorkerRow,)).where(WorkerRow.worker_id == worker_id)
            worker = db.execute(q).fetchone()
            if not worker:
                raise Http404()
            if worker.authorization not in (
                WorkerAuthorization.PENDING,
                WorkerAuthorization.REJECTED,
            ):
                raise Http409("invalid Scanning Engine state")
            worker = dict(worker)
        client = ScanApiClient(
            worker["endpoint"],
            worker["scanning_app"],
            api_key=worker["api_key"],
            request_timeout=5,
            ssl_verify=worker["ssl_verify"],
        )
        wa_dat = open(
            os.path.join(settings["storage.license"], "wa_data.dat"), "rb"
        ).read()
        wa_dat_b64 = base64.b64encode(wa_dat).decode("utf-8")
        try:
            activation_result = client.post_request(
                "/worker/add_licence", data=dict(activation_data=wa_dat_b64)
            )
        except AuthException:
            with Connection(self.current_shard_name) as (db):
                q = (
                    WorkersTable.update(
                        values=dict(authorization=WorkerAuthorization.DETACHED)
                    )
                ).where(WorkerRow.worker_id == worker["worker_id"])
                db.execute(q)
            raise Http409("Scanning Engine detached")
        except HTTP404:
            raise Http409("invalid Scanning Engine response (404)")
        except HTTP500:
            raise Http409("invalid Scanning Engine response (500)")
        except HTTPException:
            raise Http409("invalid Scanning Engine response")
        except RequestError:
            raise Http409("could not connect to the Scanning Engine")

        status = activation_result.get("status")
        max_scans = activation_result.get("empty_slots")
        current_scans = activation_result.get("current_scans")
        if activation_result.get("license_hash") != license_key_hash:
            raise Http409("invalid Scanning Engine license")
        with Connection(self.current_shard_name) as (db):
            q = (
                WorkersTable.update(
                    values=dict(
                        authorization=WorkerAuthorization.AUTHORIZED,
                        status=WorkerStatus.ONLINE,
                    )
                )
            ).where(WorkerRow.worker_id == worker["worker_id"])
            db.execute(q)
        return ClsWorkerExtended(
            scanning_app=worker["scanning_app"],
            endpoint=worker["endpoint"],
            description=worker["description"],
            worker_id=worker["worker_id"],
            status=status,
            current_scans=current_scans,
            max_scans=max_scans,
            authorization=activation_result["authorization"],
        )

    def reject_worker(self, worker_id, **_):
        if not is_setup_type_on_premise_master():
            raise Http404()
        self.license.require(Features.MULTI_ENGINE)
        require_system_user(self.request_user)
        if worker_id == LocalWorker.WORKER_ID:
            raise Http409("local worker operations not allowed")
        with Connection(self.current_shard_name) as (db):
            q = sql_select((WorkerRow,)).where(WorkerRow.worker_id == worker_id)
            worker = db.execute(q).fetchone()
            if not worker:
                raise Http404()
            if worker.authorization == WorkerAuthorization.PENDING:
                pass
            else:
                if worker.authorization == WorkerAuthorization.AUTHORIZED:
                    pass
                else:
                    raise Http409("invalid worker state")
                q = (
                    sql_select((TargetConfigurationRow.target_id,))
                    .where(TargetConfigurationRow.name == "workers")
                    .where(
                        TargetConfigurationRow.value[worker.scanning_app].astext
                        == worker_id
                    )
                )
                target_id_list = []
                for r in db.execute(q).fetchall():
                    target_id_list.append(r.target_id)

                if target_id_list:
                    raise Http409("Scanning engine is used", details=target_id_list)
                q = (
                    WorkersTable.update(
                        values=dict(
                            authorization=WorkerAuthorization.REJECTED,
                            status=WorkerStatus.DETACHED,
                        )
                    )
                ).where(WorkerRow.worker_id == worker_id)
                db.execute(q)
        should_notify_worker = False
        if should_notify_worker:
            try:
                client = ScanApiClient(
                    worker["endpoint"],
                    worker["scanning_app"],
                    api_key=worker["api_key"],
                    request_timeout=5,
                    ssl_verify=worker["ssl_verify"],
                )
                client.delete_request("/worker")
            except HTTPException as e:
                logging.exception("detaching worker failed with %s", e.result)
            except RequestError:
                logging.exception("detaching worker failed with request error")
            except Exception as e:
                logging.exception("detaching worker failed with %s", e)

            return Http204()

    def check_worker(self, worker_id, **_):
        if not is_setup_type_on_premise_master():
            raise Http404()
        self.license.require(Features.MULTI_ENGINE)
        require_system_user(self.request_user)
        if worker_id == LocalWorker.WORKER_ID:
            raise Http409("local worker operations not allowed")
        with Connection(self.current_shard_name) as (db):
            q = sql_select((WorkerRow,)).where(WorkerRow.worker_id == worker_id)
            worker = db.execute(q).fetchone()
            if not worker:
                raise Http404()
            worker = dict(worker)
        client = ScanApiClient(
            worker["endpoint"],
            worker["scanning_app"],
            api_key=worker["api_key"],
            request_timeout=5,
            ssl_verify=worker["ssl_verify"],
        )
        max_scans = None
        current_scans = None
        try:
            check_result = client.get_worker_status()
        except AuthException:
            status = "error"
            status_extra = "authentication error"
        except HTTP404:
            status = "error"
            status_extra = "invalid worker response (404)"
        except HTTP500 as e:
            status = "error"
            status_extra = "internal error in worker"
            logger.error("internal error in worker %s", e)
        except HTTPException as e:
            status = "error"
            status_extra = e.result
        except RequestError:
            status = "offline"
            status_extra = "could not connect worker"
        else:
            status = check_result.get("status")
            status_extra = None
            max_scans = check_result.get("max_scans")
            current_scans = check_result.get("current_scans")

        return ClsWorkerExtended(
            scanning_app=worker["scanning_app"],
            endpoint=worker["endpoint"],
            description=worker["description"],
            worker_id=worker["worker_id"],
            status=status,
            status_extra=status_extra,
            current_scans=current_scans,
            max_scans=max_scans,
            authorization=worker["authorization"],
        )
